import os
from stable_baselines3.common.callbacks import BaseCallback

class SaveCallback(BaseCallback):
    """
    times: 每times次保存一次模型
    printTimes: 每printTimes次打印一次
    total_timesteps： 模型训练的次数
    save_path：保存模型的路径
    """
    def __init__(self, verbose=0, times=500000, printTimes=1000, total_timesteps = 1000,save_path='./save_model'):
        super(SaveCallback, self).__init__(verbose)

        self.times = times
        self.save_path = save_path
        self.printTimes = printTimes
        self.total_timesteps = total_timesteps
        # Those variables will be accessible in the callback
        # (they are defined in the base class)
        # The RL model
        # self.model = None  # type: BaseAlgorithm
        # An alias for self.model.get_env(), the environment used for training
        # self.training_env = None  # type: Union[gym.Env, VecEnv, None]
        # Number of time the callback was called
        # self.n_calls = 0  # type: int
        # self.num_timesteps = 0  # type: int
        # local and global variables
        # self.locals = None  # type: Dict[str, Any]
        # self.globals = None  # type: Dict[str, Any]
        # The logger object, used to report things in the terminal
        # self.logger = None  # stable_baselines3.common.logger
        # # Sometimes, for event callback, it is useful
        # # to have access to the parent object
        # self.parent = None  # type: Optional[BaseCallback]

    def _on_step(self) -> bool:
        """
        This method will be called by the model after each call to `env.step()`.

        For child callback (of an `EventCallback`), this will be called
        when the event is triggered.

        :return: (bool) If the callback returns False, training is aborted early.
        """

        if(self.n_calls % self.printTimes == 0):
            loading_number = float(self.n_calls) / float(self.total_timesteps) * 100
            print("当前进度："+ str(loading_number))
            print("当前探索率: "+ str(self.model.exploration_rate))


        if(self.n_calls % self.times == 0):
            save_path = os.path.join(self.save_path, 'model_'+str(self.n_calls))
            self.model.save(save_path)

            print("阶段性保存成功-----"+str(self.n_calls))

        return True